#include <gtest/gtest.h>
#include <string>
#include <iostream>
#include <fstream>

#include "model/root.h"
#include "safetensors_loader.h"

const std::string deepseek_vl2_root = "../3rd/deepseek-vl2-tiny";

using namespace iangels;

TEST(SafetensorsLoaderTest, LoadHeaderTest) {
    std::string filename = deepseek_vl2_root + "/model-00001-of-000001.safetensors";
    std::ifstream file(filename, std::ifstream::binary);

    SafeTensorsLoader loader;
    
    SafeTensorsLoader::SafeTensor tensor;
    Ret ret = loader.load_header_size(file, tensor);
    EXPECT_EQ(ret, Ret::SUCCESS);
    EXPECT_EQ(tensor.header_bytes_, 331320);

    ret = loader.load_header(file, tensor);
    EXPECT_EQ(ret, Ret::SUCCESS);

    Config& header = tensor.header_;
    EXPECT_EQ(header["__metadata__"]["format"], "pt");
}

TEST(SafetensorsLoaderTest, MmapTensorsTest) {
    std::string filename = deepseek_vl2_root + "/model-00001-of-000001.safetensors";

    SafeTensorsLoader::SafeTensor tensor;
    tensor.file_bytes_ = std::filesystem::file_size(filename);
    tensor.header_bytes_ = 331320;

    SafeTensorsLoader loader;
    Ret ret = loader.mmap_weights(filename, tensor);
    EXPECT_EQ(ret, Ret::SUCCESS);
}

TEST(SafetensorsLoaderTest, loadTest) {
    std::string filename = deepseek_vl2_root + "/model-00001-of-000001.safetensors";

    std::vector<std::string> filenames = {filename};
    SafeTensorsLoader loader;
    Ret ret = loader.load(filenames);
    EXPECT_EQ(ret, Ret::SUCCESS);

    auto& tensors = loader.tensors();
    EXPECT_EQ(tensors.size(), 1);

    ret = loader.unmap_weights();
    EXPECT_EQ(ret, Ret::SUCCESS);
}